import os
from PIL import Image
from concurrent.futures import ThreadPoolExecutor
import multiprocessing


def compress_image(input_file, output_file, quality=85):
    """压缩单个图像文件"""
    try:
        with Image.open(input_file) as img:
            file_format = img.format
            
            try:
                img.save(output_file, 
                        format=file_format, 
                        optimize=True,
                        quality=quality)
            except OSError:
                img.save(output_file, 
                        format=file_format, 
                        optimize=True, 
                        exif=b'',
                        quality=quality)

        print(f"图像已成功压缩: {output_file}")
    except Exception as e:
        print(f"处理失败: {input_file} - 原因: {str(e)}")


def compress_images_in_folders(input_path, extensions, quality, thread_count):
    """批量压缩图像文件（支持文件/文件夹）"""
    supported_formats = ['JPEG', 'PNG', 'BMP', 'GIF', 'TIFF']
    image_files = []
    
    # 判断输入类型
    if os.path.isdir(input_path):  # 文件夹模式
        print(f"正在处理文件夹: {input_path}")
        for root, _, files in os.walk(input_path):
            for file in files:
                input_file = os.path.join(root, file)
                _, ext = os.path.splitext(input_file)
                
                if ext.lower() in extensions:
                    try:
                        with Image.open(input_file) as img:
                            if img.format in supported_formats:
                                image_files.append(input_file)
                    except Exception as e:
                        print(f"跳过无效文件: {input_file} - 错误: {str(e)}")
    
    elif os.path.isfile(input_path):  # 单文件模式
        print(f"正在处理单个文件: {input_path}")
        input_file = input_path
        _, ext = os.path.splitext(input_file)
        
        if ext.lower() in extensions:
            try:
                with Image.open(input_file) as img:
                    if img.format in supported_formats:
                        image_files.append(input_file)
            except Exception as e:
                print(f"跳过无效文件: {input_file} - 错误: {str(e)}")
        else:
            print(f"不支持的文件扩展名: {ext}")
    
    else:
        print(f"错误：路径不存在或无效 - {input_path}")
        return

    # 并行压缩处理
    if image_files:
        with ThreadPoolExecutor(max_workers=thread_count) as executor:
            for input_file in image_files:
                # 生成输出文件名（带冲突检测）
                base_name, ext = os.path.splitext(input_file)
                output_file = f"{base_name}_compressed{ext}"
                
                if os.path.exists(output_file):
                    i = 1
                    while os.path.exists(output_file):
                        output_file = f"{base_name}_compressed_{i}{ext}"
                        i += 1

                executor.submit(compress_image, input_file, output_file, quality)
    else:
        print("未找到有效图像文件进行处理")


if __name__ == "__main__":
    # 配置参数（统一管理）
    input_path = input("请输入路径（文件夹或单个图片文件）：")
    VALID_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff']
    COMPRESSION_QUALITY = 85  # 压缩质量（1-100，数值越小压缩率越高）
    
    thread_count = multiprocessing.cpu_count()
    print(f"使用 {thread_count} 个线程进行图像压缩")
    
    compress_images_in_folders(input_path, VALID_EXTENSIONS, COMPRESSION_QUALITY, thread_count)